import math

import matplotlib
import networkx as nx
import matplotlib.pyplot as plt
import json
from networkx.readwrite import json_graph
import random
import time

# 修正中文显示问题
matplotlib.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 中文字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 处理负号问题


def custom_node_attr(G, node, edges_per_new_node):
    degree = len(list(G.neighbors(node)))
    neighbor_count = degree
    centrality = degree / (2 * len(G.nodes()) - 2) if len(G.nodes()) > 1 else 0
    node_type = 'center' if degree > edges_per_new_node else 'leaf'
    max_ports = 1000
    service_count = int(centrality * max_ports)
    service_count = max(1, min(service_count, max_ports))
    port_service_pairs = random.sample(top_services, min(service_count, len(top_services)))
    ports = list({pair['port_protocol'] for pair in port_service_pairs})
    services = list({pair['service_name'] for pair in port_service_pairs if pair['service_name'] != 'unknown'})
    return {
        'type': node_type,
        'degree': degree,
        'neighbor_count': neighbor_count,
        'centrality': centrality,
        'ports': ports,
        'service': services,
    }


def test_P_ports(shared_ports, vul_probabilities):
    non_vulnerable_prob = 1.0
    for port in shared_ports:
        port_number = int(port.split('/')[0])
        vulnerability_prob = vul_probabilities.get(port_number, 0)
        non_vulnerable_prob *= (1 - vulnerability_prob)
    return 1 - non_vulnerable_prob


def custom_edge_attr(u, v, u_f, v_f, vul_p):
    centrality_sum = u_f.get('centrality', 0) + v_f.get('centrality', 0)
    centrality_diff = abs(u_f.get('centrality', 0) - v_f.get('centrality', 0))
    degree_sum = u_f.get('degree', 0) + v_f.get('degree', 0)
    weight = 1 + centrality_sum * 10 + degree_sum * 0.1
    shared_ports = list(set(u_f.get('ports', [])) & set(v_f.get('ports', [])))
    threshold = 10
    P_conn = 0.5 + min(len(shared_ports) / threshold, 1) * 0.5
    P_ports = test_P_ports(shared_ports, vul_p)
    P_uv = P_ports * P_conn
    return {
        'weight': weight,
        'centrality_sum': centrality_sum,
        'centrality_diff': centrality_diff,
        'degree_sum': degree_sum,
        'connection_type': 'connect' if u_f.get('type') == 'center' or v_f.get('type') == 'center' else 'normal',
        'vulnerability_transmission': P_uv
    }


class HierarchicalNetwork:
    def __init__(self, num_nodes, edges_per_new_node,
                 node_attr_func=None, edge_attr_func=None, model='BA', **kwargs):
        """
        :param num_nodes: 节点数
        :param edges_per_new_node: 每个新节点连接的边数（适用于 BA 模型）
        :param node_attr_func: 节点属性生成函数
        :param edge_attr_func: 边属性生成函数
        :param model: 图生成模型类型('BA', 'ER', 'WS', 'REGULAR','RGEOM')
        :param kwargs: 其他模型参数
        """
        self.num_nodes = num_nodes
        self.edges_per_new_node = edges_per_new_node
        self.node_attr_func = node_attr_func or self.default_node_attr
        self.edge_attr_func = edge_attr_func or self.default_edge_attr
        self.model = model
        self.kwargs = kwargs  # 存储其他模型参数
        self.G = self.generate_hierarchical_topology()
        self._compute_centralities()

    def default_node_attr(self, G, node, edges_per_new_node):
        return {'type': 'center' if len(list(G.neighbors(node))) > edges_per_new_node else 'leaf'}

    def default_edge_attr(self, u, v):
        return {'weight': 1}

    def test_P_vul(self):
        port_range = (1, 65532)
        random.seed(time.time())
        return {port: random.uniform(0.01, 0.5) for port in range(*port_range)}

    def generate_hierarchical_topology(self):
        #G = nx.barabasi_albert_graph(self.num_nodes, self.edges_per_new_node)
        """
        根据指定的模型生成图拓扑。
        支持的模型：'BA'（Barabási–Albert 模型）、'ER'（Erdős–Rényi 随机图）、
        'WS'（Watts–Strogatz 小世界模型）、'REGULAR'（随机正则图）、'RGEOM'（随机几何图）
        """
        if self.model == 'BA':
            # Barabási–Albert 模型
            G = nx.barabasi_albert_graph(self.num_nodes, self.edges_per_new_node)
        elif self.model == 'ER':
            # Erdős–Rényi 随机图
            edge_prob = self.kwargs.get('edge_prob', 0.1)  # 默认边生成概率为 0.1
            G = nx.erdos_renyi_graph(self.num_nodes, edge_prob)
        elif self.model == 'WS':
            # Watts–Strogatz 小世界模型
            neighbors = self.kwargs.get('neighbors', 4)  # 每个节点连接的最近邻居数
            rewire_prob = self.kwargs.get('rewire_prob', 0.1)  # 重新连边的概率
            G = nx.watts_strogatz_graph(self.num_nodes, neighbors, rewire_prob)
        elif self.model == 'REGULAR':
            # 随机正则图
            degree = self.kwargs.get('degree', 5)  # 每个节点的度数
            G = nx.random_regular_graph(degree, self.num_nodes)
        elif self.model == 'RGEOM':
            # 随机几何图
            radius = self.kwargs.get('radius', 0.2)
            G = nx.random_geometric_graph(self.num_nodes, radius)
        else:
            raise ValueError(f"Unsupported model: {self.model}")
        
        
        for node in G.nodes():
            G.nodes[node]['node_feature'] = self.node_attr_func(G, node, self.edges_per_new_node)
        vul_P = self.test_P_vul()
        for u, v in G.edges():
            G.edges[u, v].update(
                self.edge_attr_func(u, v,
                                    G.nodes[u]['node_feature'],
                                    G.nodes[v]['node_feature'],
                                    vul_P)
            )
        return G

    def _compute_centralities(self):
        self.pagerank = nx.pagerank(self.G, alpha=0.8)
        self.betweenness = nx.betweenness_centrality(self.G)
        self.closeness = nx.closeness_centrality(self.G)
        self.degree_cent = nx.degree_centrality(self.G)
        self.eigenvector = nx.eigenvector_centrality(self.G, max_iter=500)
        self.triangles = nx.triangles(self.G)
        self.clustering = nx.clustering(self.G)
        self.katz = nx.katz_centrality(self.G, alpha=0.1, beta=1.0)

    def zoom(self, levels):
        """
        根据指定的层级数合并叶子节点。
        :param levels: 需要执行的合并层级数
        """
        if levels is None or levels <= 0:
            return
        for _ in range(levels):
            leaves = [n for n, d in self.G.nodes(data=True)
                      if d['node_feature'].get('type') == 'leaf']
            if not leaves:
                break
            for leaf in leaves:
                neighbors = list(self.G.neighbors(leaf))
                if not neighbors:
                    continue
                # 选择度数最高的父节点
                parent = max(neighbors, key=lambda x: self.G.degree(x))
                parent_feat = self.G.nodes[parent]['node_feature']
                leaf_feat = self.G.nodes[leaf]['node_feature']
                for k, v in leaf_feat.items():
                    if isinstance(v, (int, float)):
                        parent_feat[k] = parent_feat.get(k, 0) + v
                # 列表特征保留重复项
                for k, v in leaf_feat.items():
                    if isinstance(v, list):
                        parent_feat[k] = parent_feat.get(k, []) + v  # 去掉了set去重
                self.G.remove_node(leaf)
            for node in self.G.nodes():
                self.G.nodes[node]['node_feature']['type'] = (
                    'center' if self.G.degree(node) > self.edges_per_new_node else 'leaf'
                )
        self._compute_centralities()
    def visualize_topology(self, title="Barabási–Albert层次结构随机网络拓扑图"):
        pos = nx.spring_layout(self.G, k=0.3, iterations=50)
        colors = ['red' if self.G.nodes[n]['node_feature'].get('type') == 'center' else 'lightblue'
                  for n in self.G.nodes()]
        nx.draw(self.G, pos, node_color=colors, with_labels=True, node_size=200)
        plt.title(title)
        plt.show(block=True)

    def export_topology_data(self, file_name="hierarchical_graph_data.json"):
        for node in self.G.nodes():
            feat = self.G.nodes[node]['node_feature']
            feat.update({
                'pagerank': self.pagerank.get(node, 0),
                'betweenness_centrality': self.betweenness.get(node, 0),
                'closeness_centrality': self.closeness.get(node, 0),
                'degree_centrality': self.degree_cent.get(node, 0),
                'eigenvector_centrality': self.eigenvector.get(node, 0),
                'triangles': self.triangles.get(node, 0),
                'clustering': self.clustering.get(node, 0),
                'katz_centrality': self.katz.get(node, 0),
            })
        bridges = set(nx.bridges(self.G))
        for u, v in self.G.edges():
            self.G[u][v]['is_bridge'] = ((u, v) in bridges or (v, u) in bridges)
        data = json_graph.node_link_data(self.G)
        with open(file_name, 'w') as f:
            json.dump(data, f, indent=4)
        print(f"拓扑数据已导出到 {file_name}")


def extract_top_k_services(k=100):
    with open('./DataProcess/nmap_services.json', 'r') as json_file:
        services = json.load(json_file)
    sorted_services = sorted(services, key=lambda x: float(x['open_frequency']), reverse=True)
    top_k_services = sorted_services[:k]
    return top_k_services

top_services = extract_top_k_services()


# 原generate函数
def generate(num_nodes, edges_per_new_node,
             gen_data_name, zoom_levels=None, visual_enable=False):
    network = HierarchicalNetwork(num_nodes, edges_per_new_node,
                                  node_attr_func=custom_node_attr,
                                  edge_attr_func=custom_edge_attr)
    if zoom_levels is not None:
        network.zoom(zoom_levels)
    if visual_enable:
        network.visualize_topology()
    network.export_topology_data(f"{gen_data_name}")

# 新增模型选择参数的generate函数
def generate(num_nodes, edges_per_new_node, gen_data_name, zoom_levels=None, visual_enable=False, model='BA', **kwargs):
    """
    生成网络拓扑。
    :param num_nodes: 节点数
    :param edges_per_new_node: 每个新节点连接的边数（适用于 BA 模型）
    :param gen_data_name: 导出的文件名
    :param zoom_levels: 层级合并数
    :param visual_enable: 是否可视化
    :param model: 图生成模型类型（'BA', 'ER', 'WS', 'REGULAR','RGEOM'）
    :param kwargs: 其他模型参数
    """
    network = HierarchicalNetwork(num_nodes, edges_per_new_node,
                                  node_attr_func=custom_node_attr,
                                  edge_attr_func=custom_edge_attr,
                                  model=model, **kwargs)
    if zoom_levels is not None:
        network.zoom(zoom_levels)
    if visual_enable:
        network.visualize_topology()
    network.export_topology_data(f"{gen_data_name}")



if __name__ == "__main__":
    generate(128, 1, 'zoomed_graph', zoom_levels=1, visual_enable=True)
